import matplotlib.pyplot as plt
import pickle
import torch
from torch.nn import functional as F
from sklearn.metrics import roc_curve, auc
import matplotlib.pyplot as plt

# from config import train_data_path
# from config import graphs_path


def load_valid(train=True):
    if train:
        with open("./train_submodel_result/ClosestMemberModel.pkl", "rb")as data:
            ClosestMemberModel = pickle.load(data)
        with open("./train_submodel_result/HistoryModel.pkl", "rb")as data:
            HistoryModel = pickle.load(data)
        with open("./train_submodel_result/OtherMemberModel.pkl", "rb")as data:
            OtherMemberModel = pickle.load(data)
        with open("./train_submodel_result/YesTendencyModel.pkl", "rb")as data:
            YesTendencyModel = pickle.load(data)
    else:
        with open("./test_submodel_result/test_submodel_result/ClosestMemberModel_test.pkl", "rb")as data:
            ClosestMemberModel = pickle.load(data)
        with open("./test_submodel_result/test_submodel_result/HistoryModel_test.pkl", "rb")as data:
            HistoryModel = pickle.load(data)
        with open("./test_submodel_result/test_submodel_result/OtherMemberModel_test.pkl", "rb")as data:
            OtherMemberModel = pickle.load(data)
        with open("./test_submodel_result/test_submodel_result/YesTendencyModel_test.pkl", "rb")as data:
            YesTendencyModel = pickle.load(data)
    
    return ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel

def loss(phi, ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel):
    real_phi = F.softmax(phi, dim=0)
    n = HistoryModel.shape[0]
    pred = real_phi[0] * ClosestMemberModel[:,1] + real_phi[1] * HistoryModel[:,1] + real_phi[2] * OtherMemberModel[:,1] + real_phi[3] * YesTendencyModel[:,1]
    pred = pred.reshape(n,1)
    target = HistoryModel[:,0].reshape(n,1)

    loss = F.binary_cross_entropy(pred, target)
    return loss

def train_gradient_descent(phi, loss, lr, ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel):
    phi.requires_grad = True
    optimizer = torch.optim.SGD(params=[phi], lr=lr) 

    for _ in range(10):
        optimizer.zero_grad()
        output = loss(phi, ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel)
        output.backward() 
        optimizer.step() 
    
    real_phi = F.softmax(phi, dim=0)
    pred = real_phi[0] * ClosestMemberModel[:,1] + real_phi[1] * HistoryModel[:,1] + real_phi[2] * OtherMemberModel[:,1] + real_phi[3] * YesTendencyModel[:,1]
    target = HistoryModel[:,0]

    fpr, tpr, thresholds = roc_curve(target.detach().numpy(), pred.detach().numpy(), pos_label=1)
    auc_value = auc(fpr, tpr)
    return real_phi, auc_value, thresholds, pred, target

def to_torch(Model, train=True):
    if train :
        reslut = torch.ones(len(Model), 2)
        for iter, item in enumerate(Model) :
            if(item[0]=="yes"):
                reslut[iter][0] = 1
                reslut[iter][1] = item[1]
            else:
                reslut[iter][0] = 0
                reslut[iter][1] = item[1]
    else:
        reslut = torch.ones(len(Model), 2)
        for iter, item in enumerate(Model) :
            if(item[2]=="yes"):
                reslut[iter][0] = 1
                reslut[iter][1] = item[3]
            else:
                reslut[iter][0] = 0
                reslut[iter][1] = item[3]
    return reslut


def find_best_threshold(thresholds, pred, target):
    best_accuracy = 0
    best_threshold = 0
    pred_temp = torch.zeros_like(pred)
    for iter in range(thresholds.shape[0]):
        pred_temp[pred>thresholds[iter]] = 1
        pred_temp[pred<=thresholds[iter]] = 0
        accuracy = (pred_temp==target).sum() / pred_temp.shape[0]
        if(accuracy > best_accuracy):
            best_accuracy = accuracy
            best_threshold = thresholds[iter]
    return best_accuracy,best_threshold

def test(real_phi, best_threshold, ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel):
    pred = real_phi[0] * ClosestMemberModel[:,1] + real_phi[1] * HistoryModel[:,1] + real_phi[2] * OtherMemberModel[:,1] + real_phi[3] * YesTendencyModel[:,1]
    target = ClosestMemberModel[:,0]
    fpr, tpr, thresholds = roc_curve(target.detach().numpy(), pred.detach().numpy(), pos_label=1)
    
    pred[pred>best_threshold] = 1
    pred[pred<=best_threshold] = 0

    accuracy =  (pred==target).sum()/pred.shape[0]
    auc_value = auc(fpr,tpr)
    plt.plot(fpr, tpr, label="ROC")
    plt.xlabel("FPR")
    plt.ylabel("TPR")
    plt.plot([0, 1], [0, 1], color='navy', linestyle='--')
    plt.savefig('fix.jpg', dpi=300) 
    return auc_value, accuracy, pred, target


if __name__ == "__main__":

    isTrain = True
    isTest = True
    if isTrain:
        ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel =load_valid(train=True)
        ClosestMemberModel = to_torch(ClosestMemberModel, train=True)
        HistoryModel = to_torch(HistoryModel, train=True)
        OtherMemberModel = to_torch(OtherMemberModel, train=True)
        YesTendencyModel  = to_torch(YesTendencyModel, train=True)

        phi = torch.rand(4)
        real_phi, auc_value,thresholds, pred, target= train_gradient_descent(phi, loss, 1e-8, ClosestMemberModel, HistoryModel, OtherMemberModel, YesTendencyModel)
        best_accuracy,best_threshold = find_best_threshold(thresholds, pred, target)

        print("#Train# ROC_AUC:{}, Accuracy:{}, Threshold:{} ".format(auc_value, best_accuracy, best_threshold))
    
    if isTest:
        ClosestMemberModel_test, HistoryModel_test, OtherMemberModel_test, YesTendencyModel_test =load_valid(train=False)
        ClosestMemberModel_test = to_torch(ClosestMemberModel_test, train=False)
        HistoryModel_test = to_torch(HistoryModel_test, train=False)
        OtherMemberModel_test = to_torch(OtherMemberModel_test, train=False)
        YesTendencyModel_test  = to_torch(YesTendencyModel_test, train=False)
        auc, accuracy, pred, target = test(real_phi, best_threshold, ClosestMemberModel_test, HistoryModel_test, OtherMemberModel_test, YesTendencyModel_test)

        print("#Test# ROC_AUC:{}, Accuracy:{}".format(auc, accuracy))